import json
import logging

from pymongo import MongoClient

fmt = '[%(asctime)s]{%(filename)s:%(lineno)d}[%(levelname)s] %(message)s'
logging.basicConfig(level=logging.DEBUG, format=fmt)
logger = logging.getLogger(__name__)

client = MongoClient('mongodb://lxt:aa347347@173.82.107.218:27017/?authSource=admin')
db = client['dns']
collection_history = db['history2']
collection_records = db['resolvers']
collection_abnormal = db['abnormal']


def check_history():
    name_ips_map = {}

    for doc in collection_history.find():
        name = doc['domain']
        ips = set()
        records = doc['dnsData']['data']['history']['a']['data']['records']
        for record in records:
            for value in record['values']:
                ips.add(value['ip'])
        name_ips_map[name] = ips

    for name, ips_real in name_ips_map.items():
        pipeline = [
            {
                '$match': {
                    'domain': f'{name}.'
                }
            }, {
                '$unwind': {
                    'path': '$A'
                }
            }, {
                '$group': {
                    '_id': {
                        'ip': '$ip',
                        'domain': '$domain'
                    },
                    'ips': {
                        '$addToSet': '$A'
                    },
                    'first_seen': {
                        '$min': '$updateTime'
                    },
                    'last_seen': {
                        '$max': '$updateTime'
                    }
                }
            }, {
                '$project': {
                    '_id': 0,
                    'ip': '$_id.ip',
                    'domain': '$_id.domain',
                    'first_seen': '$first_seen',
                    'last_seen': '$last_seen',
                    'ips': '$ips'
                }
            }
        ]

        for doc in collection_records.aggregate(pipeline):
            ip = doc['ip']

            ips = set(doc['ips'])

            abnormal_ips = ips - ips_real

            if abnormal_ips:
                logger.debug(f'[{name}][{ip}]: {abnormal_ips} - {ips_real}')
            else:
                logger.debug(f'[{name}][{ip}]: not abnormal - {ips_real}')
            collection_abnormal.update_one({'domain': name + '.', 'ip': ip}, {'$set': {
                **doc, 'abnormal_ips': list(abnormal_ips)
            }}, upsert=True)


def summarize():
    # ips_count = len(collection_records.distinct('ip'))
    aggregate = [
        {
            '$group': {
                '_id': {
                    'abnormal_ips': '$abnormal_ips',
                    'name': '$name'
                },
                'count': {
                    '$sum': 1
                },
                'resolver_ips': {
                    '$addToSet': '$ip'
                }
            }
        }
    ]
    cursor = collection_abnormal.aggregate(aggregate)
    with open('result/summarize.json', 'w') as f:
        json.dump(list(cursor), f, indent=4, ensure_ascii=False)


if __name__ == '__main__':
    check_history()
    summarize()
